import math

import numpy as np
from scipy.integrate import quad
from scipy.special import j0


# 定义设备类
class Node:
    def __init__(self, location, bandwidth, t_power, c_power=None):
        self.location = np.array(location)
        self.bandwidth = bandwidth  # 带宽
        self.t_power = t_power  # 传输功率
        self.c_power = c_power  # 计算功率（可选）
        self.p_k = 0.1
        self.f = 2e9
class User(Node):
    def __init__(self, location, bandwidth, t_power, data_size, comp_density):
        super().__init__(location, bandwidth, t_power)
        self.data_size = data_size  # 数据量
        self.comp_density = comp_density  # 计算密度


class UAV(Node):
    def __init__(self, location, bandwidth, t_power, c_power, comp_resource, p_power):
        super().__init__(location, bandwidth, t_power, c_power)
        self.comp_resource = comp_resource  # 计算资源
        self.p_power = p_power  # 推进功率
        self.velocity = 0
        # 定义其余需要的属性
        self.P0 = 79.86  # 基础功率，单位瓦特
        self.U_tip = 120  # 旋翼尖速度
        self.Pi = 88.63  # 飞行功率，单位瓦特
        self.rho_0 = 1.225  # 空气密度，单位 kg/m^3
        self.d0 = 0.6  # 机身阻力比
        self.A = 0.503  # 旋翼桨叶面
        self.S = 0.05  # 旋翼密实度
        self.V0 = 4.03  # 悬停状态下的平均旋翼感应速度，单位可以根据实际情况设置（如 m/s）
        self.H = 50
        self.gamma = 10e-27  #电容系数
    def calculate_propulsion_power(self):
        """计算并更新无人机的推进功耗"""

        # 根据公式计算推进能耗
        propulsion_energy = (self.P0 * (1 + (3 * self.velocity ** 2) / (self.U_tip ** 2)) +
                             self.Pi * np.sqrt(
                    np.sqrt(1 + (self.velocity ** 4) / (4 * self.V0 ** 2)) - (self.velocity ** 2) / (
                                2 * self.V0 ** 2)) +
                             0.5 * self.d0 * self.A * self.rho_0 * self.S * self.velocity ** 3)

        return propulsion_energy  # 计算飞行功耗

    def calculate_computation_delay(self, user):
        """计算无人机的计算时延"""
        return user.data_size * user.comp_density / self.comp_resource

    def calculate_computation_energy(self, computation):
        """计算无人机的计算能耗"""
        return computation * (self.comp_resource ** 2) * self.gamma

class BaseStation(Node):
    def __init__(self, location, comp_resource):
        super().__init__(location, bandwidth=1, t_power=0.5)
        self.comp_resource = comp_resource  # 基站的计算资源
        self.H = 50
    def calculate_computation_delay(self, user):
        """计算基站的计算时延"""
        return user.data_size * user.comp_density / self.comp_resource

class Train(Node):
    def __init__(self, location, comp_resource):
        super().__init__(location, bandwidth=1, t_power=0.4)
        self.comp_resource = comp_resource  # 列车的计算资源
        self.H = 50

    def calculate_computation_delay(self, user):
        """计算列车的计算时延"""
        return user.data_size * user.comp_density / self.comp_resource


def calculate_channel_gain(node_1, node_2, a, b, mu_LOS, mu_NLoS):
    """
    参数：
    SN_node：信号节点对象
    other_node：另一信号节点对象（UAV 或 Train）
    fc：float，上行链路信道频率（Hz）
    mu_LOS：float，LoS信道的额外衰减因子
    mu_NLoS：float，NLoS信道的额外衰减因子
    """
    # 计算节点间的三维距离
    d_km = np.linalg.norm(node_1.location - node_2.location)
    d_km = math.sqrt(d_km ** 2 + node_2.H ** 2)  # 加上垂直高度
    # 计算仰角 θ_km
    theta_km = (180 / math.pi) * math.asin(node_2.H / d_km)
    # 计算 LoS 概率 p_LOS_km 和 NLoS 概率 p_NLoS_km
    p_LOS_km = 1 / (1 + a * math.exp(-b * (theta_km - a)))
    p_NLoS_km = 1 - p_LOS_km
    # 计算信道损耗 PL_LOS_km 和 PL_NLoS_km
    c = 3e8  # 光速 (m/s)
    PL_LOS_km = 20 * math.log10((4 * math.pi * node_1.f * d_km) / c) + mu_LOS
    PL_NLoS_km = 20 * math.log10((4 * math.pi * node_1.f * d_km) / c) + mu_NLoS
    # 计算信道增益 g_km
    g_km = (p_LOS_km * 10 ** (-PL_LOS_km / 10)) + (p_NLoS_km * 10 ** (-PL_NLoS_km / 10))
    return g_km

def UtoB_data_rate(node_1, node_2, a, b, mu_LOS, mu_NLoS, SIGMA_2):
    g_km = calculate_channel_gain(node_1, node_2, a, b, mu_LOS, mu_NLoS)
    # 计算信干噪比 γ_km
    gamma_km = (node_1.p_k * g_km) / SIGMA_2
    # 计算数据传输速率 R_km
    R_km = node_1.bandwidth * math.log2(1 + gamma_km)
    return R_km

def data_rate(node_1, node_2, a, b, mu_LOS, mu_NLoS, SIGMA_2):
    g_km = calculate_channel_gain(node_1, node_2, a, b, mu_LOS, mu_NLoS)
    # 计算信干噪比 γ_km
    gamma_km = (node_1.p_k * g_km) / SIGMA_2
    # 计算数据传输速率 R_km
    R_km = node_1.bandwidth * math.log2(1 + gamma_km)
    return R_km


def UtoT_data_rate(node_1, Train_node, a, b, mu_LOS, mu_NLoS, SIGMA_2, I_ici):
    g_km = calculate_channel_gain(node_1, Train_node, a, b, mu_LOS, mu_NLoS)
    print(g_km)
    # 计算信干噪比 γ_km
    gamma_km = (node_1.p_k * g_km) / SIGMA_2 + g_km * I_ici
    # 计算数据传输速率 R_km
    R_km = node_1.bandwidth * math.log2(1 + gamma_km)
    return R_km


def doppler_interference_factor(fm_T, vT, c, Ts):
    # 被积函数定义
    def integrand(tau):
        term = 2 * np.pi * fm_T * (-vT / c) * Ts  # 参数计算
        return 1 - abs(tau) * j0(term)  # 被积函数 1 - |τ| * J_0(...)
    # 执行积分，积分范围为 [-1, 1]
    integral_result = quad(integrand, -1, 1)[0]
    # 计算多普勒干扰因子
    I_ICI = 1 - integral_result
    return I_ICI


# 性能模拟
def simulate_performance():
    # 参数设置
    a = 9.61
    b = 0.16
    mu_LOS = 1
    mu_NLOS = 20
    SIGMA_2 = 1e-11  # 噪声功率
    fm_T = 5.8e9
    vT = 50
    c = 3e8
    Ts = 1
    # 初始化用户、无人机、基站、列车的位置和设备参数
    user = User(location=[100, 30, 0], bandwidth=1e6, t_power=0.1, data_size=2e6, comp_density=1e3)
    uav = UAV(location=[50., 50., 50.], bandwidth=1e7, t_power=0.5, c_power=0.5, comp_resource=1.5e9, p_power=0.3)
    bs = BaseStation(location=[200, 40, 0], comp_resource=5e9)
    train = Train(location=[350, 50, 0], comp_resource=2.5e9)
    print(f"无人机的推进能耗: {uav.calculate_propulsion_power() * 1:.2f} ")
    # 计算用户到无人机的传输性能
    rate_user_uav = data_rate(user, uav, a, b, mu_LOS, mu_NLOS, SIGMA_2)
    upload_time_user_uav = user.data_size / rate_user_uav
    user_transmission_energy = upload_time_user_uav * user.t_power
    print(f"用户到无人机的传输速率: {rate_user_uav / 1e6:.2f} Mbps")
    print(f"用户到无人机的传输时延: {upload_time_user_uav:.2f} 秒")
    print(f"用户到无人机的传输能耗: {user_transmission_energy:.2f} J")

    # 计算无人机到基站的传输性能
    rate_uav_bs = data_rate(uav, bs, a, b, mu_LOS, mu_NLOS, SIGMA_2)
    upload_time_uav_bs = user.data_size / rate_uav_bs
    uav_transmission_energy_bs = upload_time_uav_bs * uav.t_power
    print(f"无人机到基站的传输速率: {rate_uav_bs / 1e6:.2f} Mbps")
    print(f"无人机到基站的传输时延: {upload_time_uav_bs:.2f} 秒")
    print(f"无人机到基站的传输能耗: {uav_transmission_energy_bs:.2f} J")

    # 计算无人机到列车的传输性能
    I_ICI = doppler_interference_factor(fm_T, vT, c, Ts)
    rate_uav_train = UtoT_data_rate(uav, train, a, b, mu_LOS, mu_NLOS, SIGMA_2, I_ICI)
    upload_time_uav_train = user.data_size / rate_uav_train
    uav_transmission_energy_train = upload_time_uav_train * uav.t_power
    print(f"无人机到列车的传输速率: {rate_uav_train / 1e6:.2f} Mbps")
    print(f"无人机到列车的传输时延: {upload_time_uav_train:.2f} 秒")
    print(f"无人机到列车的传输能耗: {uav_transmission_energy_train:.2f} J")
    # 计算无人机的计算时延和能耗
    uav_computation_delay = uav.calculate_computation_delay(user)
    uav_computation_energy = uav.calculate_computation_energy(user.data_size * user.comp_density)
    print(f"无人机的计算时延: {uav_computation_delay:.4f} 秒")
    print(f"无人机的计算能耗: {uav_computation_energy:.4f} J")
    # 计算基站的计算时延
    bs_computation_delay = bs.calculate_computation_delay(user)
    print(f"基站的计算时延: {bs_computation_delay:.4f} 秒")
    train_computation_delay = train.calculate_computation_delay(user)
    print(f"列车的计算时延: {train_computation_delay:.4f} 秒")
# 运行模拟
if __name__ == "__main__":
    simulate_performance()
